import torch
import torchvision
from torch import nn
from torch.nn import L1Loss, MSELoss, Sequential, Conv2d, MaxPool2d, Flatten, Linear
from torch.utils.data import DataLoader


# inputs=torch.tensor([1,2,3],dtype=torch.float)
# targets=torch.tensor([1,2,5],dtype=torch.float)
#
#
# inputs=torch.reshape(inputs,(1,1,1,3))
# targets=torch.reshape(targets,(1,1,1,3))
#
# loss=L1Loss(reduction='sum')
# result=loss(inputs,targets)
#
# loss_mse=MSELoss()
# result_mes=loss_mse(inputs,targets)
#
# print(result_mes)
#
# x=torch.tensor([0.1,0.2,0.3])
# y=torch.tensor([1])
# x=torch.reshape(x,(1,3))
# loss_cross=nn.CrossEntropyLoss()
# result_cross=loss_cross(x,y)
# print(result_cross)

class JJw(nn.Module):
    def __init__(self):
        super(JJw,self).__init__()
        self.model1=Sequential(
            Conv2d(3, 32, 5, padding=2, stride=1),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self,x):
        x=self.model1(x)
        return x

dataset=torchvision.datasets.CIFAR10("./dataset",train=False,transform=torchvision.transforms.ToTensor(),download=True)
dataloader=DataLoader(dataset,batch_size=1)

loss=nn.CrossEntropyLoss()
jjw=JJw()
optim = torch.optim.SGD(jjw.parameters(),lr=0.01)
for data in dataloader:
    imgs,targets=data
    outputs=jjw(imgs)
    result_loss=loss(outputs,targets)
    result_loss.backward()
    optim.zero_grad()
    optim.step()



